# 功能：对（选择题，阅读题，翻译题，作文题）实现ocr,并将其返回为正确的格式
import requests
import base64
import re
import os
import json
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage
import sys

SPARKAI_URL = 'wss://spark-api.xf-yun.com/v3.1/chat'
SPARKAI_APP_ID = '1d38b17a'
SPARKAI_API_SECRET = 'MTJjOGQ5YzE1ZWZjYjJkOWUxZTM4Mzcx'
SPARKAI_API_KEY = 'b33d686049c006fd12a7e9b345ec0409'
SPARKAI_DOMAIN = 'generalv3'

access_token = "24.9825cfb6262e262a5caa6fef4f531037.2592000.1713240501.282335-56894520"
api_key = "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"

spark = ChatSparkLLM(
    spark_api_url=SPARKAI_URL,
    spark_app_id=SPARKAI_APP_ID,
    spark_api_key=SPARKAI_API_KEY,
    spark_api_secret=SPARKAI_API_SECRET,
    spark_llm_domain=SPARKAI_DOMAIN,
    streaming=False,
)

def ocr(baidu_api, url, type):
    """
    实现对不同题型拍照
    :param baidu_api: 百度ocr_key
    :param src: 图片的路径
    :param type: ocr类型
    :return: 对应的json格式，不同题，不同json
    """
    access_token = baidu_api
    request_url = "https://aip.baidubce.com/rest/2.0/ocr/v1/accurate"
    f = open(url, "rb")
    img = base64.b64encode(f.read())
    params = {"image": img, "language_type": "CHN_ENG", "result_type": "big"}
    request_url = request_url + "?access_token=" + access_token
    headers = {"content-type": "application/x-www-form-urlencoded"}
    response = requests.post(request_url, data=params, headers=headers)

    if type == "yuedu":
        max_width = 0
        for words in response.json()["words_result"]:
            max_width = max(max_width, words["location"]["width"])
        article = " "
        for words in response.json()["words_result"]:
            if words["location"]["width"] >= max_width - 50:
                article = article + words["words"] + " "
            else:
                article = article + words["words"] + "\n" + "  "
        return article
    if type == "xuanze":
        max_width = 0
        for words in response.json()["words_result"]:
            max_width = max(max_width, words["location"]["width"])
        choice_partten = "^[abcdABCD]\..*"
        article = " "
        now_words = response.json()["words_result"][0]
        same_row = 0
        for i, words in enumerate(response.json()["words_result"]):
            if i > 0:
                if (
                    now_words["location"]["top"] + 10
                    > words["location"]["top"]
                    > now_words["location"]["top"] - 10
                ):
                    same_row = same_row + 1
                    if re.match(choice_partten, words["words"]):
                        article = article + now_words["words"] + " "
                        same_row = 0
                    else:
                        article = article + now_words["words"] + "( )"
                else:
                    if same_row == 0 and not re.match(choice_partten, now_words["words"]):
                        article = article + now_words["words"] + "( )"
                    else:
                        article = article + now_words["words"] + " "
                now_words = words
        article = article + response.json()["words_result"][-1]["words"]
        return article
    if type == "fanyi" or type == "zuowen":
        max_width = 0
        for words in response.json()["words_result"]:
            max_width = max(max_width, words["location"]["width"])
        article = " "
        for words in response.json()["words_result"]:
            if words["location"]["width"] >= max_width - 50:
                article = article + words["words"] + " "
            else:
                article = article + words["words"] + "\n" + "  "
        return article

def call_sparkai(prompt):
    messages = [ChatMessage(role="user", content=prompt)]
    handler = ChunkPrintHandler()
    response = spark.generate([messages], callbacks=[handler])
    return response

def final_ocr_res(baidu_api, spark_key, url, type):
    os.environ["http_proxy"] = "http://127.0.0.1:10808"
    os.environ["https_proxy"] = "http://127.0.0.1:10808"
    ocr_res = ocr(baidu_api, url, type)
    if type == "yuedu":
        reading_ocr_prompt = """
        请从给定的英语阅读理解问题中生成一个JSON对象，其中包括原文、问题和选项、以及每个问题所考察的题型(例如主旨大意题、细节理解题、推理判断题、词汇理解题等等)。JSON应该为原始文本和每个选项提供单独的字段。这是阅读理解问题：确保原文的格式与输入中的格式完全相同，并且每个选项在JSON对象中都被分隔到自己的字段中。输出应采用有效的JSON格式。下面是输入和输出的一个格式
        输入：一篇英文阅读，包括原文和选项
        输出：只能以json格式输出,下面这个json就是对应的输出
        {
        'original_text': 'This is the original text of the reading comprehension question. It should be formatted exactly as it appears in the input.',
        'questions': [
            {
            'question_number': 1,
            'question': 'the reading comprehension question',
            'options': [
                {'Option': 'This is the option for the reading comprehension question.'},
                ] ,
            'question_type': 'The type of this question',
            }
        ]
        }
        要求：
        （1）现在根据上面的提示，对下面这篇阅读题进行分割，并且返回对应的json格式。只返回json格式，不能有其他废话。
        （2）在original_text对应的字符串中，必须保持原来输入的阅读段落格式，不能合并成一段话，这个很重要，必须完成！！！也就是必须保证阅读正文对应的换行符位置不能改变！！！
        （3）在question_type应该是该问题所属题型，如主旨大意题、细节理解题、推理判断题、词汇理解题，句子理解题，文章结构题,作者观点题,文体风格题等等。其值必须用中文表示！
        (5)在json格式中的值，如果出现' " \\等可能会导致报错的符号，需要使用转义符来来将其转义，保证json格式的正确性
        输入的英语阅读如下：
            {ocr_res}"""
        result = call_sparkai(reading_ocr_prompt)
        res_reading = json.loads(result)
        return_data = {"category": "yuedu"}
        original_text = res_reading["original_text"]
        sonExe = []
        for item in res_reading["questions"]:
            reading_exe = {}
            options = item["options"]
            size = len(options)
            reading_exe["choiceNum"] = size
            reading_exe["problem"] = item["question"]
            for i in range(size):
                reading_exe["choice" + str(i + 1)] = options[i]["Option"]
            reading_exe["question_type"] = item["question_type"]
            sonExe.append(reading_exe)
        return_data["contenxt"] = original_text
        return_data["sonExercise"] = sonExe
        return_data["stuAnswer"] = ""
        return_data["field"] = ""
        return_data["aiComment"] = ""
        print(str(return_data))

    if type == "xuanze":
        choice_ocr_prompt_part1 = """
        现在需要你将一串包含多道英语选择题的字符串（提示：可能含有其他话，你需要将其删除）转换成易于阅读的格式，并以 JSON 的形式返回。在 JSON 中，你需要包含原问题数组，每个原问题中不仅包含原文，还要包含选项。
        首先，你需要明白字符串原本的格式可能非常混乱，每个题目之间以及每个选项之间仅有一个空格隔开，没有换行符。并且，可能会含有一些与题目无关的一些话，你也需要将其删除。因此，你的任务是将其划分成正确的格式，使其适合人们阅读，并且将其转换成 JSON 格式。
        JSON 格式的示例如下：
        {
          "questions": [
            {
              "question": "What is the capital of France?",
               "options": [
                {"Option": "This is the first option for the reading comprehension question."},
                {"Option": "This is the second option for the reading comprehension question."},
                ]
            },
            {
              "question": "Who wrote 'Romeo and Juliet'?",
              "options": [
                {"Option": "This is the first option for the reading comprehension question."},
                {"Option": "This is the second option for the reading comprehension question."},
                ]
            }
          ]
        }
        在这个示例中，`questions` 键对应一个数组，数组中每个元素都是一个对象，代表一道题目。每个题目对象中有两个键：`question` 表示题目内容，`options` 表示选项，是一个包含多个选项字符串的数组。
        因此，你需要按照这个格式将原始字符串中的题目和选项提取出来，组织成 JSON 格式返回。这样做可以使得信息更清晰易读，也更便于程序处理和解析。
        下面是输入的英语选择题字符串：
        """
        choice_ocr_prompt_part2 = """
        注意：
        只能以示例中的json格式返回，不能输出除json外任何的单词
        """
        final_prompt = choice_ocr_prompt_part1 + ocr_res + choice_ocr_prompt_part2
        res = call_sparkai(final_prompt)
        res_reading = json.loads(res)
        return_data = {"category": "xuanze"}
        options = res_reading[0]["options"]
        return_data["problem"] = res_reading[0]["question"]
        for item in options:
            size = len(options)
            return_data["choiceNum"] = size
            for i in range(size):
                return_data["choice" + str(i + 1)] = options[i]["Option"]
        print(str(return_data))

    if type == "zuowen" or type == "fanyi":
        print(ocr_res)

if __name__ == "__main__":
    url = sys.argv[1]
    url = os.path.abspath(".") + "/api/target/classes/static/photo/" + url
    type = sys.argv[2]
    final_ocr_res(baidu_api=access_token, spark_key=api_key, url=url, type=type)
